package com.hsu.hsurpc.registry;

import cn.hutool.core.collection.ConcurrentHashSet;
import com.hsu.hsurpc.config.RegistryConfig;
import com.hsu.hsurpc.model.ServiceMetaInfo;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.CuratorCache;
import org.apache.curator.framework.recipes.cache.CuratorCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;

import java.time.LocalDateTime;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * ZooKeeper 注册中心实现类，使用 Curator 框架
 * @Author Hsu琛君珩
 * @Date 2024-09-22 13:11
 * @Description
 * @Version: v1.0.0
 */
@Slf4j
public class ZooKeeperRegistry implements Registry {

    /**
     * ZooKeeper 客户端，用于与 ZooKeeper 服务器交互
     */
    private CuratorFramework client;

    /**
     * 服务发现组件
     */
    private ServiceDiscovery<ServiceMetaInfo> serviceDiscovery;

    /**
     * 本地注册的服务节点集合
     */
    private final Set<String> localRegistryNodeKeySet = new HashSet<>();

    /**
     * 注册中心服务缓存（支持多个服务键）
     */
    private final RegistryServiceMultiCache registryServiceMultiCache = new RegistryServiceMultiCache();

    /**
     * 正在监听的 key 集合
     */
    private final Set<String> watchingKeySet = new ConcurrentHashSet<>();

    /**
     * ZooKeeper 根节点路径，防止不同项目的数据冲突
     */
    private static final String ZK_ROOT_PATH = "/rpc/zk";

    /**
     * 初始化 ZooKeeper 客户端，连接到注册中心
     *
     * @param registryConfig 注册中心的配置信息，包含连接地址、超时时间等
     */
    @Override
    public void init(RegistryConfig registryConfig) {
        client = CuratorFrameworkFactory.builder()
                .connectString(registryConfig.getAddress())                                                           // 使用配置的地址
                .retryPolicy(new ExponentialBackoffRetry(Math.toIntExact(registryConfig.getTimeout()), 3)) // 设置重试策略
                .build();
        serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceMetaInfo.class)
                .client(client)
                .basePath(ZK_ROOT_PATH)
                .serializer(new JsonInstanceSerializer<>(ServiceMetaInfo.class))
                .build();
        try {
            client.start();             // 启动 ZooKeeper 客户端
            serviceDiscovery.start();   // 启动服务发现组件
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 注册服务到 ZooKeeper 注册中心
     *
     * @param serviceMetaInfo 服务的元数据信息（包括服务名称、版本、地址、端口等）
     * @throws Exception 如果注册失败，抛出异常
     */
    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws Exception {
        // 设置注册时间
        serviceMetaInfo.setRegisterTime(LocalDateTime.now());
        serviceDiscovery.registerService(builderServiceInstance(serviceMetaInfo));
        String registerKey = ZK_ROOT_PATH + "/" + serviceMetaInfo.getServiceNodeKey();
        localRegistryNodeKeySet.add(registerKey); // 添加节点信息到本地缓存
    }

    /**
     * 注销服务节点
     * 从 ZooKeeper 注册中心中删除已注册的服务节点信息，并从本地缓存中移除该节点的键名。
     *
     * @param serviceMetaInfo 要注销的服务的元数据信息
     */
    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) {
        try {
            serviceDiscovery.unregisterService(builderServiceInstance(serviceMetaInfo));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        String registerKey = ZK_ROOT_PATH + "/" + serviceMetaInfo.getServiceNodeKey();
        localRegistryNodeKeySet.remove(registerKey); // 从本地缓存移除
    }

    /**
     * 服务发现，根据服务键名前缀，查找已注册的所有服务节点信息
     *
     * @param serviceKey 服务的唯一键名（通常由服务名称和版本号组成）
     * @return 返回该服务下的所有节点列表
     */
    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
        List<ServiceMetaInfo> cachedServiceMetaInfoList = registryServiceMultiCache.readCache(serviceKey);
        if (cachedServiceMetaInfoList != null) {
            return cachedServiceMetaInfoList;
        }
        try {
            Collection<ServiceInstance<ServiceMetaInfo>> serviceInstanceList = serviceDiscovery.queryForInstances(serviceKey);
            List<ServiceMetaInfo> serviceMetaInfoList = serviceInstanceList.stream()
                    .map(ServiceInstance::getPayload)
                    .collect(Collectors.toList());
            serviceMetaInfoList.forEach(service -> {
                System.out.println("Service registered at: " + service.getRegisterTime());
            });
            registryServiceMultiCache.writeCache(serviceKey, serviceMetaInfoList); // 写入缓存
            return serviceMetaInfoList;
        } catch (Exception e) {
            throw new RuntimeException("获取服务列表失败", e);
        }
    }

    /**
     * 启动心跳机制
     * 对于 ZooKeeper 的临时节点不需要额外的心跳机制，
     * 因为当服务器故障时，临时节点会自动被删除。
     */
    @Override
    public void heartBeat() {
        // ZooKeeper 的临时节点不需要心跳机制，节点会在服务器故障时自动删除
    }

    /**
     * 监听指定服务节点的变化
     * 通过 ZooKeeper 的 Watch 机制，监控指定服务节点的状态变化（如删除事件），
     * 当服务节点发生变化时，执行相应的处理逻辑。
     *
     * @param serviceNodeKey 需要监听的服务节点的键名
     */
    @Override
    public void watch(String serviceNodeKey) {
        String watchKey = ZK_ROOT_PATH + "/" + serviceNodeKey;
        boolean newWatch = watchingKeySet.add(watchKey);
        if (newWatch) {
            CuratorCache curatorCache = CuratorCache.build(client, watchKey);
            curatorCache.start();
            curatorCache.listenable().addListener(
                    CuratorCacheListener.builder()
                            .forDeletes(childData -> registryServiceMultiCache.clearCache(serviceNodeKey))          // 监听删除事件，清理缓存
                            .forChanges(((oldNode, node) -> registryServiceMultiCache.clearCache(serviceNodeKey)))  // 监听变化事件，清理缓存
                            .build()
            );
        }
    }

    /**
     * 销毁注册中心，释放 ZooKeeper 客户端资源
     * 在项目关闭或服务停止时调用此方法，断开与 ZooKeeper 注册中心的连接，
     * 下线所有已注册的服务节点，并释放相关的客户端资源。
     */
    @Override
    public void destroy() {
        log.info("当前节点下线");
        for (String key : localRegistryNodeKeySet) {
            try {
                client.delete().guaranteed().forPath(key);
            } catch (Exception e) {
                throw new RuntimeException(key + "节点下线失败", e);
            }
        }
        if (client != null) {
            client.close(); // 关闭 ZooKeeper 客户端连接
        }
    }

    private ServiceInstance<ServiceMetaInfo> builderServiceInstance(ServiceMetaInfo serviceMetaInfo) {
        String serviceAddress = serviceMetaInfo.getServiceHost() + ":" + serviceMetaInfo.getServicePort();
        try {
            return ServiceInstance.<ServiceMetaInfo>builder()
                    .id(serviceAddress)
                    .name(serviceMetaInfo.getServiceKey())
                    .address(serviceAddress)
                    .payload(serviceMetaInfo)
                    .build();
        } catch (Exception e) {
            throw new RuntimeException("构建服务实例失败", e);
        }
    }
}
